import torch
import torchvision
from PIL import Image
from torch import nn

image_path = "images/frog.png"
image = Image.open(image_path)
image = image.convert("RGB")
print(image)

# image.show()

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),
    torchvision.transforms.ToTensor()
])
image = transform(image)
print(image.shape)


# 搭建神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.model(x)
        return x


model = torch.load("net_0.pth")
# print(model)

image = torch.reshape(image, (1, 3, 32, 32))
model.eval()  # 切换为评估模式
with torch.no_grad():  # 在该上下文管理器内，禁用梯度计算：
    output = model(image)
print(output)
print(torch.argmax(output))
# 映射成中文
class_list = ["飞机", "汽车", "鸟", "猫", "鹿", "狗", "青蛙", "马", "船", "卡车"]
print(class_list[torch.argmax(output)])
